import json
import os
import sys
from typing import List

import pandas as pd


def merge_strategy_table(strategy_file_list: List[str], dir_output: str) -> None:
    """
    将多个生成的预测文件拼接，如果对同一个股票和同一个时刻有多个预测结果，则进行平均
    :param strategy_file_list: 策略文件
    :param dir_output: 输出策略结果到该文件夹中
    :return:
    """
    os.makedirs(dir_output, exist_ok=True)
    merge_info = dict(strategy_file_list=strategy_file_list)
    with open(os.path.join(dir_output, "merge_info.json"), "w") as f:
        json.dump(merge_info, f, ensure_ascii=False)

    file_result = os.path.join(dir_output, "merge_pred.csv")

    dfs = []
    for f in strategy_file_list:
        df = pd.read_csv(f, index_col=0)
        dfs.append(df)

    all_df = pd.concat(dfs, axis=0)
    all_df = all_df.reset_index(drop="date_time" in all_df.columns)
    merge_df = all_df.groupby(["date_time", "time", "code_market"]).mean()
    merge_df.to_csv(file_result)
    print(f"total merge: {strategy_file_list}")


if __name__ == "__main__":

    partsdir, filename, outdir = sys.argv[1:]
    subdirs = os.listdir(partsdir)
    strategy_dirs = [os.path.join(partsdir, x) for x in subdirs]
    strategy_file_list = [os.path.join(x, filename) for x in strategy_dirs]
    merge_strategy_table(strategy_file_list, outdir)
